import numpy as np
from matplotlib import pyplot as plt
from sklearn import preprocessing
from sklearn.cross_decomposition import PLSRegression
from sklearn.preprocessing import PolynomialFeatures
import getData
from LS_minPYX import train

degree = 10
poly = PolynomialFeatures(degree=degree)


class CDM():
    def fit(self, data):
        data.getData()
        sigma = 0.5
        lbd_reg = 1e-4
        lbd_inv = 0.1
        lng_rate = 0.1
        max_iter = 100
        Thresh = 1e-5

        # cl=====================================================
        Xs_ = data.source_x_train.reshape(-1)
        Ys = data.source_y_train[:, 0][:, np.newaxis]
        Xt_ = data.target_x_train.reshape(-1)
        Yt = data.target_y_train[:, 0][:, np.newaxis]
        # Yt = data.target_y_train[:, 1][:, np.newaxis]

        Xs = np.vstack((Xs_, Xs_)).T
        Xt = np.vstack((Xt_, Xt_)).T

        Ys_new = train(Xs, Xt, Ys, Yt, sigma=sigma,
                       lambda_regularization=lbd_reg,
                       lambda_inv=lbd_inv,
                       learning_rate=lng_rate,
                       Max_Iter=max_iter,
                       Thresh=Thresh)

        self.Ys_new = Ys_new.detach().numpy()

        x_train = data.source_x_train
        x_train = preprocessing.scale(x_train)
        x_train = poly.fit_transform(x_train)
        self.cdm_model = pls_model(x_train, self.Ys_new)

        # cd=============================================================
        Xs_ = data.source_x_train.reshape(-1)
        Ys = data.source_y_train[:, 1][:, np.newaxis]
        Xt_ = data.target_x_train.reshape(-1)
        Yt = data.target_y_train[:, 1][:, np.newaxis]
        # Yt = data.target_y_train[:, 1][:, np.newaxis]

        Xs = np.vstack((Xs_, Xs_)).T
        Xt = np.vstack((Xt_, Xt_)).T

        Ys_new = train(Xs, Xt, Ys, Yt, sigma=sigma,
                       lambda_regularization=lbd_reg,
                       lambda_inv=lbd_inv,
                       learning_rate=lng_rate,
                       Max_Iter=max_iter,
                       Thresh=Thresh)

        self.Ys_new = Ys_new.detach().numpy()

        x_train = data.source_x_train
        x_train = preprocessing.scale(x_train)
        x_train = poly.fit_transform(x_train)
        self.cdm_model1 = pls_model(x_train, self.Ys_new)

    def predict(self, x_test):
        x_test = preprocessing.scale(x_test)
        x_test = poly.fit_transform(x_test)
        self.y_pre = self.cdm_model.predict(x_test)
        self.y_pre1 = self.cdm_model1.predict(x_test)
        self.y_pre_all = np.hstack((self.y_pre, self.y_pre1))

        return self.y_pre_all


def pls_model(x_train_st, y_train):
    n_components = 5
    pls = PLSRegression(n_components=n_components)
    pls.fit(x_train_st, y_train)
    return pls


if __name__ == '__main__':
    cdm = CDM()
    data = getData.GetData()
    model = cdm.fit(data)
    x = data.target_x_test

    # x_st = preprocessing.scale(x)
    # x_st = poly.fit_transform(x_st)
    # y_pre = model.predict(x_st)

    y_pre = cdm.predict(x)

    y = data.target_y_test
    for i in range(y_pre.shape[1]):
        plt.figure()
        plt.plot(x, y_pre[:, i], color="red", label="pre")
        # plt.plot(data.source_x_train, cdm.Ys_new[:, i], label='Y_new')
        plt.plot(x, y[:, i], label="true")
        plt.legend(loc=0, frameon=False)
        plt.show()
